/* Copyright (c) 2022, TD SYNNEX Corporation. All rights reserved */

package com.synnex.base.jpa.repository;

import jakarta.persistence.EntityManager;
import lombok.extern.slf4j.Slf4j;
import org.hibernate.StaleStateException;
import org.hibernate.engine.internal.Versioning;
import org.hibernate.engine.jdbc.spi.JdbcCoordinator;
import org.hibernate.engine.spi.SharedSessionContractImplementor;
import org.hibernate.jdbc.BatchFailedException;
import org.hibernate.jdbc.BatchedTooManyRowsAffectedException;
import org.hibernate.persister.entity.AbstractEntityPersister;
import org.hibernate.persister.entity.EntityPersister;
import org.hibernate.resource.jdbc.spi.LogicalConnectionImplementor;
import org.hibernate.sql.exec.spi.JdbcOperation;
import org.hibernate.sql.model.MutationOperation;
import org.hibernate.sql.model.MutationOperationGroup;
import org.hibernate.type.Type;
import org.springframework.orm.jpa.EntityManagerFactoryUtils;
import org.springframework.orm.jpa.EntityManagerProxy;
import org.springframework.util.Assert;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Collection;
import java.util.List;



@Slf4j
public class RepositoryUtil {

    /**
     * @param entities
     * @param updateSQL
     * @param entityManager
     * @param entityPersister
     * @return
     * @throws SQLException
     */
    public static <T> int executeBatchUpdate(Collection<? extends T> entities, String updateSQL, EntityManager entityManager,
                                             EntityPersister entityPersister, List<List<Object>> paramsList) throws SQLException {

        EntityManager transactionalEntityManager = EntityManagerFactoryUtils.doGetTransactionalEntityManager(entityManager.getEntityManagerFactory(), null, true);
        SharedSessionContractImplementor session = getSharedSessionContractImplementor(transactionalEntityManager);
        // Connection connection1 = session.connection();
        JdbcCoordinator jdbcCoordinator = session.getJdbcCoordinator();
        LogicalConnectionImplementor logicalConnection = jdbcCoordinator.getLogicalConnection();
        Connection physicalConnection = logicalConnection.getPhysicalConnection();
        try (PreparedStatement prepareStatement = session.getJdbcCoordinator().getStatementPreparer().prepareStatement(updateSQL)) {
            int i = 0;
            for (T entity : entities) {
                int index = 1;
                List<Object> properties = paramsList.get(i++);
                if (entityPersister.isVersioned()) {
                    int lastIdx = properties.size() - 1;
                    Object curVersion = properties.get(lastIdx);
                    Object increment = Versioning.incrementVersion(entity, curVersion, entityPersister, session);
                    properties.set(lastIdx, increment);
                }
                for (Object param : properties) {
                    prepareStatement.setObject(index++, param);
                }
                Object identifier = entityPersister.getIdentifier(entity, session);
                Type idType = entityPersister.getIdentifierType();
                idType.nullSafeSet(prepareStatement, identifier, index, session);
                index += getIdentifierColumnSpan(entityPersister);
                if (entityPersister.isVersioned()) {
                    Type versionType = entityPersister.getVersionType();
                    Object version = entityPersister.getVersion(entity);
                    versionType.nullSafeSet(prepareStatement, version, index, new boolean[]{true}, session);
                }
                prepareStatement.addBatch();
            }
            int[] affectedRows = prepareStatement.executeBatch();
            int sum = 0;
            for (int affectedRow : affectedRows) {
                checkBatched(affectedRow, 0, updateSQL);
                sum += affectedRow;
            }
            return sum;
        }
    }

    // full copy com.synnex.base.jpa.repository.RepositoryUtil.executeBatch,
    // only remove isAutoGenerateId. and add version auto_increment
    public static <T, ID> int executeBatchUpdate(Collection<? extends T> entities, String updateSQL, EntityManager entityManager,
                                                 EntityPersister entityPersister) throws SQLException {
        EntityManager transactionalEntityManager = EntityManagerFactoryUtils.doGetTransactionalEntityManager(entityManager.getEntityManagerFactory(), null, true);
        SharedSessionContractImplementor session = getSharedSessionContractImplementor(transactionalEntityManager);
        try (PreparedStatement prepareStatement = session.getJdbcCoordinator().getStatementPreparer().prepareStatement(updateSQL, Statement.RETURN_GENERATED_KEYS);) {
            Type[] propertyTypes = entityPersister.getPropertyTypes();
            boolean[] updateability = entityPersister.getPropertyUpdateability();
            boolean[][] columnInsertable = getPropertyColumnUpdateable(entityPersister);
            for (T entity : entities) {
                int index = 1;
                Object[] properties = entityPersister.getPropertyValuesToInsert(entity, null, null);
                if (entityPersister.isVersioned()) {
                    Object version = entityPersister.getVersion(entity);
                    Assert.notNull(version, "version field is null");
                    Object increment = Versioning.incrementVersion(entity, version, entityPersister, session);
                    Versioning.setVersion( properties, increment, entityPersister);
                }
                for (int i = 0, size = propertyTypes.length; i < size; i++) {
                    if (updateability[i]) {
                        Type propertyType = propertyTypes[i];
                        propertyType.nullSafeSet(prepareStatement, properties[i], index++, columnInsertable[i], session);
                    }
                }

                // the id(composite id) is at the end.
                Object id = entityPersister.getIdentifier(entity, session);
                Type idType = entityPersister.getIdentifierType();
                idType.nullSafeSet(prepareStatement, id, index, session);
                index += getIdentifierColumnSpan(entityPersister);
                if (entityPersister.isVersioned()) {
                    Type versionType = entityPersister.getVersionType();
                    Object version = entityPersister.getVersion(entity);
                    versionType.nullSafeSet(prepareStatement, version, index, new boolean[]{true}, session);
                }
                prepareStatement.addBatch();
            }
            int[] affectedRows = prepareStatement.executeBatch();
            int sum = 0;
            for (int affectedRow : affectedRows) {
                checkBatched(affectedRow, 0, updateSQL);
                sum += affectedRow;
            }
            return sum;
        }
    }


    public static final int expectedRowCount = 1;

    /**
     * org.hibernate.jdbc.Expectations.BasicExpectation#checkBatched(int, int, java.lang.String)
     * @param rowCount
     * @param batchPosition
     * @param statementSQL
     */
    private static void checkBatched(int rowCount, int batchPosition, String statementSQL) {
        if ( rowCount == -2 ) {
            log.debug( "Success of batch update unknown: {}", batchPosition );
        }
        else if ( rowCount == -3 ) {
            throw new BatchFailedException( "Batch update failed: " + batchPosition );
        }
        else {
            if ( expectedRowCount > rowCount ) {
                throw new StaleStateException(
                        "Batch update returned unexpected row count from update ["
                                + batchPosition + "]; actual row count: " + rowCount
                                + "; expected: " + expectedRowCount + "; statement executed: "
                                + statementSQL
                );
            }
            if ( expectedRowCount < rowCount ) {
                String msg = "Batch update returned unexpected row count from update [" +
                        batchPosition + "]; actual row count: " + rowCount +
                        "; expected: " + expectedRowCount;
                throw new BatchedTooManyRowsAffectedException( msg, expectedRowCount, rowCount, batchPosition );
            }
        }
    }

    public static int getIdentifierColumnSpan(EntityPersister entityPersister) {
        AbstractEntityPersister abstractEntityPersister = (AbstractEntityPersister) entityPersister;
        return abstractEntityPersister.getIdentifierColumnSpan();
    }
    /**
     * @param entityManager
     * @return
     * @see org.springframework.orm.jpa.ExtendedEntityManagerCreator#createProxy
     */
    static SharedSessionContractImplementor getSharedSessionContractImplementor(EntityManager entityManager) {
        // 1. no Proxy, org.hibernate.internal.SessionImpl implements SharedSessionContractImplementor and EntityManager
        if (entityManager instanceof SharedSessionContractImplementor) {
            return (SharedSessionContractImplementor) entityManager;
        }
        // 2. java.lang.reflect.Proxy, org.springframework.orm.jpa.ExtendedEntityManagerCreator#createProxy
        EntityManager iterator = entityManager;
        while (iterator instanceof EntityManagerProxy) {
            iterator = ((EntityManagerProxy) iterator).getTargetEntityManager();
            if (iterator instanceof SharedSessionContractImplementor) {
                return (SharedSessionContractImplementor) iterator;
            }
        }
        // 3. org.springframework.aop.framework.AopProxy
        // AopUtils.isJdkDynamicProxy()
        return null;
    }


    public static String getSQLUpdateStrings(EntityPersister entityPersister) {
        AbstractEntityPersister abstractEntityPersister = (AbstractEntityPersister) entityPersister;

        return extractSqlStrings(abstractEntityPersister.getUpdateCoordinator().getStaticMutationOperationGroup())[0];
        // return abstractEntityPersister.getSQLUpdateStrings()[0];
    }


    public static boolean[][] getPropertyColumnUpdateable(EntityPersister entityPersister) {
        AbstractEntityPersister abstractEntityPersister = (AbstractEntityPersister) entityPersister;
        return abstractEntityPersister.getPropertyColumnUpdateable();
    }


    private static String[] extractSqlStrings(MutationOperationGroup operationGroup) {
        final int numberOfOperations = operationGroup.getNumberOfOperations();
        final String[] strings = new String[numberOfOperations];
        for ( int i = 0; i < numberOfOperations; i++ ) {
            final MutationOperation operation = operationGroup.getOperation( i );
            if ( operation instanceof JdbcOperation) {
                strings[i] = ( (JdbcOperation) operation ).getSqlString();
            }
        }
        return strings;
    }
}
